package cz.cuni.amis.pogamut.episodic.schemas;

import cz.cuni.amis.pogamut.episodic.decisions.Action;
import cz.cuni.amis.pogamut.episodic.decisions.AtomicAction;
import cz.cuni.amis.pogamut.episodic.decisions.Intention;
import cz.cuni.amis.pogamut.episodic.decisions.Node;
import cz.cuni.amis.pogamut.episodic.decisions.NodeType;
import cz.cuni.amis.pogamut.episodic.memory.Affordance;
import cz.cuni.amis.pogamut.episodic.memory.AgentMemory;
import cz.cuni.amis.pogamut.episodic.memory.IdGenerator;
import cz.cuni.amis.pogamut.episodic.memory.Parameters;
import cz.cuni.amis.pogamut.episodic.schemas.utils.Binomial;
import cz.cuni.amis.pogamut.episodic.episodes.Episode;
import cz.cuni.amis.pogamut.episodic.episodes.EpisodeNode;
import cz.cuni.amis.pogamut.episodic.episodes.ObjectNode;
import cz.cuni.amis.pogamut.episodic.episodes.ObjectSlot;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.io.Serializable;
import java.util.*;
import org.apache.commons.collections15.multimap.MultiHashMap;

/**
 * The <code>SchemaBag</code> is one of the important classes when it comes
 * to agent's memories. It does not hold any particular episode, insteas it
 * is responsible for creating an average episodes, so called schemas. It
 * accumullates all different experiences instances of one episode into one
 * aggregated schema. Apart from the capability to retrieve universal episodes
 * it gives us a way how to tell what sub-episodes happen more often that
 * others and it helps us to more efficiently store the concrete episodes by
 * providing a way to derive some of the nodes from what is only a part of
 * complete episode tree.
 * <p>
 * Schemabag consist of <code>SchemaEpisodeNodes</code> which are equivalent to
 * <code>EpisodeNode</code>s and <code>SchemaObjectNodes</code> which are
 * equivalent to <code>ObjectNode</code>s. Counters are kept for number of
 * experiences containing particular nodes and for all set of nodes up to
 * specified size. Theses <code>SchemaCounter</code>s are kept in
 * <code>SchemaBag</code> and indexed in a way that it is always fast to find
 * all the counters representing subsets or supersets of a set of nodes
 * represented by a specified counter.
 *
 * @author Michal Cermak
 */
public class SchemaBag implements Serializable {
    /**
     * Determines if a de-serialized file is compatible with this class.
     *
     * Maintainers must change this value if and only if the new version
     * of this class is not compatible with old versions. See Sun docs
     * for <a href=http://java.sun.com/products/jdk/1.1/docs/guide
     * /serialization/spec/version.doc.html> details. </a>
     *
     * Not necessary to include in first version of the class, but
     * included here as a reminder of its importance.
     */
    private static final long serialVersionUID = 1L;

    /**
     * A reference to ID generator structure, that will provide
     * the unique IDs to all inside node. This unique IDs are mostly
     * needed when identifying the visualized vertices (because nodes
     * can have identical names).
     */
    IdGenerator idGenerator;

    /**
     * A reference to the parent structure. It should be possible
     * to access decision tree, chronobags or other parts of mem via
     * this reference.
     */
    AgentMemory mem;

    /**
     * Count the total number of  nodes in the schemabag. Used mostly to
     * determine whether there was new node added since last refresh of
     * schemabag visualization, so that it can be redrawed.
     */
    private int numberOfNodes = 0;

    /**
     * ID of the schemabag. Used as an ID of vertex representing
     * the schemabag when visualizing Chronobag View.
     */
    public final int id;

    /**
     * <code>schemaENodes</code> is a collection of all <code>SchemaEpisodicNode</code>s
     * in the <code>SchemaBag</code>. Since upon its creation each node
     * is mutually associated with a node in decision tree, these nodes
     * are in a map <strong>indexed by the unique ID of their associated
     * nodes</strong>.
     */
    public HashMap<String, SchemaEpisodeNode> schemaENodes = new HashMap<String, SchemaEpisodeNode>();

    /**
     * <code>schemaENodes</code> is a collection of all <code>SchemaObjectNode</code>s
     * in the <code>SchemaBag</code>. Since each item has to have a unique name,
     * all nodes are indexed in a map <strong>indexed by the name of an item</strong>.
     */
    public HashMap<String, SchemaObjectNode> schemaONodes = new HashMap<String, SchemaObjectNode>();

    /**
     * <code>counters</code> is a collection of all <code>SchemaCounter</code>s
     * in the <code>SchemaBag</code> inside a container. There is one global
     * counters, and K counters for each temporal schema of periodicity K.
     */
    public TemporalSchema counters = null;

    public int currentDay = 0;

    /**
     * Because counters in the schemabag are updated more often than
     * the individual nodes, it is convenient not to redraw all whole
     * schemabag visualization each time they are increased. But still,
     * the detailed information that is invoded when the count number is
     * requested by the user needs to be up to date. Therefore the reference
     * to <code>SchemaMessageCommand</code> class is passed to the
     * visualizer and it is used to receive latest count number each time
     * it is requested by the user via visualizer window.
     */
    public class SchemaMessageCommand implements ISchemaMessageCommand, Serializable {
        /**
         * Once reference to <code>SchemaMesssageCommand</code> is set
         * in visualizer, this method can be used to retrieve up to date
         * count number of currently selected set of schema nodes.
         * @param picked    Collection of IDs of schema episode nodes and
         * schema object nodes that are selected in the visualizer and should
         * form the set of nodes representing wanted counter.
         * @return  The string message to be displayed to the user. If count
         * for selected set of nodes existed in <code>SchemaBag</code> it
         * should contain the count number of this set of nodes.
         */
        public String getSchemaMessage(Collection<Integer> picked) {
            if (picked.isEmpty()) return "Select some episodic nodes, slot contents and retry.";
            int key = 0;
            for (Integer i : picked) key += i;
            // TODO: Visualisation of other levels?
            Collection<SchemaCounter> counts = counters.getSchema().get(key);
            SchemaCounter count = null;
            if (counts != null) {
                for (SchemaCounter c : counts) {
                    if (c.checkNodes(picked)) count = c;
                }
            }
            if (count == null) return "No count for given selection.";

            String s = "Count for given selection: " + count.getCount() + System.getProperty("line.separator");
            s += "Total count for given selection: " + count.getTotalCount();
            return s;
        }
    }

    /**
     * Instantion of a <code>SchemaMessageCommand</code> class that is passed
     * to visualizer so that it can recieve up to date information about
     * <code>SchemaCounter</code>s.
     * @see SchemaMessageCommand
     */
    public SchemaMessageCommand schemaMessageCommand = new SchemaMessageCommand();

    /*
     * If execution of an action continues for a longer time, it may be misleading
     * to count it as several different executions. Thus counters that should be
     * increased in last step has to be remembered so they are not increased again.
     */
    private Collection<Integer> lastIds = new HashSet<Integer>();

    /**
     * If the trace of nodes is very similar with the one when the schema was
     * last updated, it may need be wanted to increase all the counters. Only the
     * counters containing new nodes should be increased. IDs of these nodes
     * are stored in this collection.
     */
    private Collection<Integer> requiredIds = new HashSet<Integer>();

    /**
     * Instantiate the class by providing references to a common ID generator
     * and parent <code>AgentMemory</code> structure.
     * <p>
     * The ID of the schemabag is also generated in constructor.
     *
     * @param   idGen Reference to common <code>IdGenerator</code>.
     * @param   mem Reference to the <code>AgentMemory</code> structure.
     */
    public SchemaBag(IdGenerator idGen, AgentMemory mem) {
        idGenerator = idGen;
        this.mem = mem;
        id = idGenerator.getNewId();
        TemporalSchema.maxLevels = new Integer(mem.parameters.MAX_CHRONOBAG_LEVELS);
        TemporalSchema.levelLengths = mem.parameters.CHRONOBAG_INTERVAL_LENGHTS;
        counters = new TemporalSchema();
    }

    /**
     * Returns a <code>String</code> object representing the
     * <code>SchemaBag</code>'s info. It is used to provide detailed
     * information about schemabag when it is invoked from
     * the visualizion structures.
     *
     * @return  a string representation of the value of this object.
     */
    @Override
    public String toString() {
        String newline = System.getProperty("line.separator");
        String s = "";
        s += ((Integer)id).toString() + " Nodes:" + numberOfNodes + newline;
        s += "Number of Schema Episode Nodes: " + schemaENodes.size() + newline;
        s += "Number of Schema Object Nodes: " + schemaONodes.size() + newline;
        s += "Number of Global Counters: " + counters.getSchema().size() + newline;
        for (int i=1; i < mem.parameters.MAX_CHRONOBAG_LEVELS-1; i++) {
            for (int j=0; j < mem.parameters.CHRONOBAG_INTERVAL_LENGHTS[i]; j++) {
                s += "Number of Counters in level " + i + ", periodicity " + j + ": " + counters.getSchema().size() + newline;
            }
        }
        return s;
    }

    /**
     * Getter method for the <code>numberOfNodes</code> variable.
     *
     * @return     Returns the count the total number of  nodes in the
     * schemabag. Used mostly to determine whether there was new node added
     * since last refresh of schemabag visualization, so that it can be redrawed.
     */
    public int getNumberOfNodes() {
        return numberOfNodes;
    }

    /**
     * Method used to return a single <code>SchemaEpisodicNode</code>.
     * Schema episodic nodes are indexed by IDs of associated decision nodes.
     *
     * @return  Returns a <code>SchemaEpisodicNode</code> specified by
     * the id in paremeter.
     */
  //  public SchemaEpisodeNode getSchemaENode(int id) {
   //     return schemaENodes.get(id);
   // }
    public SchemaEpisodeNode getSchemaENode(String name) {
        return schemaENodes.get(name);
    }
    /**
     * Basically a getter method for the <code>schemaENodes</code> variable.
     *
     * @return  Returns a collection of all <code>SchemaEpisodicNode</code>s
     * in the <code>SchemaBag</code>.
     */
    public Collection<SchemaEpisodeNode> getSchemaENodes() {
        return schemaENodes.values();
    }

    /**
     * Basically a getter method for the <code>schemaONodes</code> variable.
     *
     * @return  Returns a collection of all <code>SchemaObjectNode</code>s
     * in the <code>SchemaBag</code>.
     */
    public Collection<SchemaObjectNode> getSchemaONodes() {
        return schemaONodes.values();
    }

    /**
     * Retrieve the current total count of collective appearences of a set
     * of specified schema nodes.
     * @param schemaNodes Set of <code>SchemaEpisodeNode</code>s of the counter.
     * @param objectNodes Set of <code>SlotContent</code>s of the counter.
     * @return Returns the total count of collective appearences of specified nodes.
     * If no counter for specified set of nodes exists in schemabag, it means
     * there is too much nodes in the set or that no counter was created because
     * the specified nodes did not appear collectively in short-term mem when
     * performing agent's action. Based on the reason the method returns either
     * null when the set of nodes can never contribute to one counter or zero
     * when the set of nodes did not happen collectively yet.
     */
    public Integer getCount(int level, int index, Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes) {
        SchemaCounter c = getCounter(level, index, schemaNodes, objectNodes, mem.parameters.MAX_SCHEMA_COMBINATION_COUNT);
        if (c == null) return 0;
        return c.getCount();

    }

    /**
     * Retrieve the counter representing total count of collective appearences of a set
     * of specified schema nodes.
     * @param schemaNodes Set of <code>SchemaEpisodeNode</code>s of the counter.
     * @param objectNodes Set of <code>SlotContent</code>s of the counter.
     * @return Returns the reference to a <code>SchemaCounter</code> representing
     * count of collective appearences of specified nodes.
     * If no counter for specified set of nodes exists in schemabag, it means
     * there is too much nodes in the set or that no counter was created because
     * the specified nodes did not appear collectively in short-term mem when
     * performing agent's action. Based on the reason the method returns either
     * null when the set of nodes can never contribute to one counter or zero
     * when the set of nodes did not happen collectively yet.
     */
    public SchemaCounter getCounter(int level, int index, Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes, int maxCombination) {
        Collection<Integer> col = new HashSet<Integer>();
        int key = 0;
        for (SchemaEpisodeNode n : schemaNodes) {
            col.add(n.getId());
            key += n.getId();
        }
        for (SlotContent s : objectNodes) {
            col.add(s.getId());
            key += s.getId();
        }
        if (col.isEmpty()) return null;
        if (col.size() > maxCombination) return null;

        Collection<SchemaCounter> counts = counters.getSchema(level, index).get(key);
        SchemaCounter count = null;
        if (counts != null) {
            for (SchemaCounter c : counts) {
                if (c.checkNodes(col)) count = c;
            }
        }
        return count;
    }

    /**
     * Increases the count of given combination of schema nodes by one.
     * Combination can consist of both schema episodic nodes and slot contents.
     * <p>
     * If no such counter exists yet, it is created, added to the set of counters
     * in the schemabag and its count number set to one.
     *
     * @param      schemaNodes  Collection of <code>SchemaEpisodicNode</code>s in the counter.
     * @param      objectNodes  Collection of <code>SlotContent</code>s in the counter.
     */
    private void increaseCount(Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes) {
        int schemaSize = schemaNodes.size() + objectNodes.size();
        if ((schemaSize == 0) || ((schemaSize != 1) && (schemaSize < mem.parameters.HIGHLEVEL_SCHEMA_MIN_SIZE))) {
            return;
        }
        int hash = 0;
        for (SchemaEpisodeNode e : schemaNodes) {
            hash += e.id;
        }
        for (SlotContent c : objectNodes) {
            hash += c.id;
        }

        if (schemaNodes.isEmpty() && objectNodes.isEmpty()) return;

        for (int i=0; i < TemporalSchema.maxLevels-1; i++) {
            Collection<SchemaCounter> counts = counters.getSchema(i, currentDay).get(hash);
            SchemaCounter counter = null;
            if (counts != null) {
                for (SchemaCounter c : counts) {
                    if (c.checkNodes(schemaNodes, objectNodes))
                    {
                        counter = c;
                        break;
                    }
                }
            }

            if (counter == null) {
                counter = new SchemaCounter(i, currentDay, hash, schemaNodes, objectNodes);
                counters.getSchema(i, currentDay).put(hash, counter);
            }
            counter.increaseTotalCount();

            boolean increase = false;
            for (Integer j : requiredIds) {
                if (counter.containsNode(j))  {
                    increase = true;
                    break;
                }
            }
            if (increase) {
                counter.increaseCount();
            }
        }

        return;
    }

    /**
     * Generates given subset of given size of specified list of objects.
     * Used to generate subsets of schema episode nodes and slot contents
     * so that all the combination counters can be increased. Repetitive
     * calling for increasing index i generates all the existing subsets of
     * specified set of nodes.
     * <p>
     * The method is generic, objects will be either of type
     * <code>SchemaEpisodeNode</code> or <code>SlotContent</code>.
     *
     * @param      set  the set that will be returned (param only exists
     * so it is not allocated each time this method runs,
     * values in the set do not matter)
     * @param      list  Array of objects forming the complete set.
     * @param      k  Specifies the size of returned subset.
     * @param      i  Index of subset of specified size to be returned.
     *
     * @return     Returns i-th subset of the set specified by list parameter.
     */
    private <T> Collection<T> getSubSet(Collection<T> set, ArrayList<T> list, int k, int i) {
        int n = list.size();
        int found = 0;
        set.clear();
        if (k > n) return set;
        long subsum;

        while (k > 0) {
            subsum = Binomial.binomial(n-1, k-1);
            int index = 0;
            while (i >= subsum) {
                index++;
                subsum += Binomial.binomial(n-1-index, k-1);
            }
            set.add(list.get(index+found));

            i -= subsum - Binomial.binomial(n-1, k-1);
            n--;
            found++;
            k--;
        }

        return set;
    }

    /**
     * Increases the count of <strong>all</strong> given combinations
     * of specified <code>SchemaEpisodeNode</code>s and <code>SlotContent</code>s
     * up to the size of combination being MAX_COMBINATION_COUNT.
     * Combinations consist of both schema episodic nodes and slot contents.
     * <p>
     * First single counters for all schema nodes are increased. If MAX_COMBINATION_COUNT
     * is greater than one, all couters for all pairs of nodes are increased
     * by one, etc.
     * <p>
     * If no such counters exists yet, they are created, added to the set of counters
     * in the schemabag and their count number set to one.
     *
     * @param      schemaNodes  Collection of <code>SchemaEpisodicNode</code>s in the set.
     * @param      objectNodes  Collection of <code>SlotContent</code>s in the set.
     * @param      maxCombination   Determines the max size of generated combinations.
     * Should be set to "Parameters.MAX_SCHEMA_COMBINATION_COUNT".
     */
    void increaseBinomialCounts(Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes, int maxCombination) {
        // NOTE: The use of HashSet here means that combinatorial schemas cannot contain multiple same-named items.
        Collection<SchemaEpisodeNode> eSubSet = new HashSet<SchemaEpisodeNode>();
        Collection<SlotContent> oSubSet = new HashSet<SlotContent>();

        // Increase combination counts.
        ArrayList<SchemaEpisodeNode> schemaNodes2 = new ArrayList<SchemaEpisodeNode>(schemaNodes);
        ArrayList<SlotContent> objectNodes2 = new ArrayList<SlotContent>(objectNodes);
        // Increase combination counts.
        for (int s1 = 0; s1 <= maxCombination; s1++) {
            long i1max = Binomial.binomial(schemaNodes2.size(), s1);
            for (int i1 = 0; i1 < i1max; i1++) {
                eSubSet = getSubSet(eSubSet, schemaNodes2, s1, i1);
                for (int s2 = 0; s1 + s2 <= maxCombination; s2++) {
                    long i2max = Binomial.binomial(objectNodes2.size(), s2);
                    for (int i2 = 0; i2 < i2max; i2++) {
                        oSubSet = getSubSet(oSubSet, objectNodes2, s2, i2);
                        if (eSubSet.size() == s1 && oSubSet.size() == s2) {
                            increaseCount(eSubSet, oSubSet);
                        }
                    }
                }
            }
        }
        increaseCounts(schemaNodes, objectNodes);
    }

    /**
     * Increases the count of given lists.
     */
    void increaseCounts(Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes) {
        // Increase combination counts.
        ArrayList<SchemaEpisodeNode> schemaNodes2 = new ArrayList<SchemaEpisodeNode>(schemaNodes);
        ArrayList<SlotContent> objectNodes2 = new ArrayList<SlotContent>(objectNodes);

        // Special case, add a schema for the whole trace with and without objects.
        increaseCount(schemaNodes2, objectNodes2);
        increaseCount(schemaNodes2, new ArrayList<SlotContent>());
    }

    public Collection<SchemaCounter> getAllExistingSubSets(int level, int index, Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes, int maxSize) {
        Collection<SchemaCounter> result = new HashSet<SchemaCounter>();

        Collection<SchemaEpisodeNode> eSubSet = new HashSet<SchemaEpisodeNode>();
        Collection<SlotContent> oSubSet = new HashSet<SlotContent>();

        ArrayList<SchemaEpisodeNode> schemaNodes2 = new ArrayList<SchemaEpisodeNode>(schemaNodes);
        ArrayList<SlotContent> objectNodes2 = new ArrayList<SlotContent>(objectNodes);

        SchemaCounter counter = null;
        for (int s1 = 0; s1 <= maxSize; s1++) {
            long i1max = Binomial.binomial(schemaNodes2.size(), s1);
            for (int i1 = 0; i1 < i1max; i1++) {
                eSubSet = getSubSet(eSubSet, schemaNodes2, s1, i1);
                for (int s2 = 0; s1 + s2 <= maxSize; s2++) {
                    long i2max = Binomial.binomial(objectNodes2.size(), s2);
                    for (int i2 = 0; i2 < i2max; i2++) {
                        oSubSet = getSubSet(oSubSet, objectNodes2, s2, i2);
                        if (eSubSet.size() == s1 && oSubSet.size() == s2) {
                            if (s1 + s2 > 0) {
                           //     assert(getCounter(eSubSet, oSubSet, maxSize) != null);
                                counter = getCounter(level, index, eSubSet, oSubSet, maxSize);
                                if (counter != null) {
                                    result.add(counter);
                                }
                            }
                        }
                    }
                }
            }
        }

        return result;
    }

    /**
     * This method locates all the SchemaEpisodeNode that were in short-term
     * mem during the execution of atomic action and it slot locates all
     * necessary SlotContent nodes for all the items that were used.
     * When both sets of nodes are aquired, <code>increaseCounts</code>
     * method is called ensuring all necessary counts will be updated.
     *
     * @param newNodeName  Name of a node that was opened
     * @param trace List of names of decision nodes that lead to
     * the atomic action executed. The top-level goal is first in the list and
     * parent of atomic action is last item in the list.
     * @param affordances   List of <code>Affordance</code> objects.
     * Each object represent one filled affordance slot that was
     * needed for the execution of current atomic action.
     */
    public void updateSchema(String newNodeName, ArrayList<String> _trace, ArrayList<Affordance> affordances) throws Exception {
        ArrayList<String> trace = new ArrayList<String>(_trace); // Duplicate to make it immutable

        HashMap<String, SchemaEpisodeNode> eNodesSubSet = new HashMap<String, SchemaEpisodeNode>();
        HashSet<SlotContent> contentsSubSet = new HashSet<SlotContent>();
        mem.getDecisionTree().ensureNodes(newNodeName, trace, affordances, idGenerator);

        ArrayList<Node> decisionNodes = mem.getDecisionTree().getTrace(trace);
        if (decisionNodes.size() == 0) {
            return;
        }

        // Fundamentally, the newNodeName is either an action or a goal.
        Node last = decisionNodes.get(decisionNodes.size() - 1);
        AtomicAction a;
        try {
            a = ((Action)last).getAtomicAction(newNodeName);
            //if newNodeName is _search... it will not be in decision tree and 'a' will be null.
            if (a != null) decisionNodes.add(a);
        } catch (Exception e) {
            // If it's not an atomic action, it must be Intention or Action. Which is fine.
            //return;
        }

        Node node;
        Affordance aff;
        for (int i = 0; i < decisionNodes.size(); i++) {
            node = decisionNodes.get(i);
            String nodeName = "";
            SchemaEpisodeNode schemaNode;
            try {
                nodeName = node.getName();
            } catch (Exception e) {
                System.err.println(e.getMessage());
                System.err.println("SchemaBag: could not retrive ID of associated node.");
                continue;
            }
            if (!schemaENodes.containsKey(nodeName)) {
                schemaNode = new SchemaEpisodeNode(node, idGenerator.getNewId(), node.getName());
                numberOfNodes++;
                schemaENodes.put(nodeName, schemaNode);
            } else {
                schemaNode = schemaENodes.get(nodeName);
                if (node.getAssociatedNode() == null) {
                    node.setAssociatedSchemaNode(schemaNode);
                }
            }
            eNodesSubSet.put(schemaNode.name, schemaNode);
        }
        for (int i = 0; i < affordances.size(); i++) {
            aff = affordances.get(i);
            if (aff.item == null) continue;
            SchemaEpisodeNode schemaNode = eNodesSubSet.get(newNodeName);
            if (schemaNode == null) continue;
            SchemaSlot slot = schemaNode.getSlot(aff.type);
            if (slot == null) {
                if (schemaNode.addSlot(aff.type, idGenerator.getNewId())) {
                    numberOfNodes++;
                }
                slot = schemaNode.getSlot(aff.type);
            }
            SchemaObjectNode object = schemaONodes.get(aff.item);
            if (object == null) {
                object = new SchemaObjectNode(idGenerator.getNewId(), aff.item);
                numberOfNodes++;
                schemaONodes.put(aff.item, object);
            }
            SlotContent slotContent = slot.getSlotContent(aff.item);
            if (slotContent == null) {
                if (slot.addSlotContent(idGenerator.getNewId(), object)) {
                    numberOfNodes++;
                }
                slotContent = slot.getSlotContent(aff.item);
            }
            contentsSubSet.add(slotContent);
        }

        requiredIds.clear();
        boolean add = false;
        for (Node n : decisionNodes) {
            if (add || schemaENodes.get(n.getName()) == null || !lastIds.contains(schemaENodes.get(n.getName()).getId())) {
                requiredIds.add(schemaENodes.get(n.getName()).getId());
                add = true;
            }
        }
        for (SlotContent c : contentsSubSet) {
            if (!lastIds.contains(c.getId())) {
                requiredIds.add(c.getId());
            }
        }
        lastIds.clear();
        for (SchemaEpisodeNode n : eNodesSubSet.values()) {
            lastIds.add(n.getId());
        }
        for (SlotContent c : contentsSubSet) {
            lastIds.add(c.getId());
        }

        increaseBinomialCounts(eNodesSubSet.values(), contentsSubSet, mem.parameters.MAX_SCHEMA_COMBINATION_COUNT);
    }

    /**
     * This method records a whole finished episode into the schema.
     *
     * When both sets of nodes are aquired, <code>increaseCounts</code>
     * method is called ensuring all necessary counts will be updated.
     *
     * @param e The episode to record.
     */
    public void updateSchema(Episode e) throws Exception {
        // Collect all nodes that currently contribute to this episode.
        Collection<SchemaEpisodeNode> schemaENodes = new HashSet<SchemaEpisodeNode>();
        Collection<SlotContent> schemaContents = new HashSet<SlotContent>();
        // Walk through the episode to collect info
        Stack<EpisodeNode> stack = new Stack<EpisodeNode>();
        SchemaSlot sSlot;
        SchemaEpisodeNode schemaNode;
        EpisodeNode stackNode;
        stack.add(e.root);
        while (!stack.isEmpty()) {
            stackNode = stack.pop();
            if (stackNode.associatedNode == null) {
                continue;
            }
            schemaNode = stackNode.associatedNode.getAssociatedNode();
            schemaENodes.add(schemaNode);
            for (ObjectSlot slot : stackNode.getObjectSlots()) {
                if (slot.getType().equals("Other")) {
                    continue;
                }
                if ((slot.getParentNode() == null) || (slot.getParentNode().associatedNode == null)) {
                    continue;
                }
                sSlot = slot.getParentNode().associatedNode.getAssociatedNode().getSlot(slot.getType());
                for (ObjectNode obj : slot.getUsedObjects()) {
                    schemaContents.add(sSlot.getSlotContent(obj.getName()));
                }
            }
            stack.addAll(stackNode.getChildrenNodes());
        }

        increaseCounts(schemaENodes, schemaContents);
    }

    public class EpisodeComparator implements Comparator<SchemaEpisodeNode> {
        private int level;
        private int index;
        public EpisodeComparator(int _level, int _index) {
            level = _level;
            index = _index;
        }
        @Override
        public int compare(SchemaEpisodeNode e1, SchemaEpisodeNode e2) {
            return e1.counts.getSchema(level, index).size() - e2.counts.getSchema(level, index).size();
        }
    }
    public class ObjectComparator implements Comparator<SlotContent> {
        private int level;
        private int index;
        public ObjectComparator(int _level, int _index) {
            level = _level;
            index = _index;
        }
        @Override
        public int compare(SlotContent e1, SlotContent e2) {
            return e1.counts.getSchema(level, index).size() - e2.counts.getSchema(level, index).size();
        }
    }

    // Filter for maximum size of a member counter.
    public HashSet<SchemaCounter> schemaCounterSubsetMaxsizeFilter(HashSet<SchemaCounter> in, int maxEpisodeSize, int maxObjectSize) {
        Iterator<SchemaCounter> it = in.iterator();
        while(it.hasNext()) {
            SchemaCounter s = it.next();
            if ((s.episodeNodes.size() > maxEpisodeSize) || (s.objectNodes.size() > maxObjectSize)) {
                it.remove();
            }
        }
        return in;
    }

    // Filter for maximum size of a member counter.
    public HashSet<SchemaCounter> schemaCounterSubsetMinsizeFilter(HashSet<SchemaCounter> in, int minEpisodeSize, int minObjectSize) {
        Iterator<SchemaCounter> it = in.iterator();
        while(it.hasNext()) {
            SchemaCounter s = it.next();
            if ((s.episodeNodes.size() < minEpisodeSize) || (s.objectNodes.size() < minObjectSize)) {
                it.remove();
            }
        }
        return in;
    }

    public class CounterComparator implements Comparator<SchemaCounter> {
        @Override
        public int compare(SchemaCounter e1, SchemaCounter e2) {
            return e2.getCount() - e1.getCount();
        }
    }

    // Filter for maximum size of a member counter.
    public ArrayList<SchemaCounter> schemaCounterSort(Collection<SchemaCounter> in) {
        ArrayList<SchemaCounter> out = new ArrayList<SchemaCounter>(in);
        Collections.sort(out, new CounterComparator());
        return out;
    }

    /**
     * This function searches the schemabag for given node cues, and returns
     * a list of matching counters containing the given nodes.
     *
     * @param schemaNodes episode node cues
     * @param objectNodes object node cues
     */
    public HashSet<SchemaCounter> schemaCounterSupersetSelector(int level, int index, Collection<SchemaEpisodeNode> schemaNodes, Collection<SlotContent> objectNodes) {
        // In order to minimize the number of list operations, let's scout the cues first. Each schemaNode has an associated
        // map of counters associated with the given node.
        ArrayList<SchemaEpisodeNode> episodeNodeList = new ArrayList<SchemaEpisodeNode>(schemaNodes);
        ArrayList<SlotContent> objectNodeList = new ArrayList<SlotContent>(objectNodes);

        // Sorting these by size will give us the optimum search. In theory, beacause the list is a multihashmap,
        // so the visible count is not exactly the size.
        Collections.sort(episodeNodeList, new EpisodeComparator(level, index));
        Collections.sort(objectNodeList, new ObjectComparator(level, index));

        return schemaCounterSupersetSelectorOne(level, index, null, episodeNodeList, objectNodeList);
    }

    /**
     * This function does a subselection on a given SchemaCounter set.
     *
     * @param current Set to be filtered further.
     * @param schemaNodes episode node cues
     * @param objectNodes object node cues
     * @return
     */
    public HashSet<SchemaCounter> schemaCounterSupersetSelectorOne(int level, int index, HashSet<SchemaCounter> current, ArrayList<SchemaEpisodeNode> schemaNodes, ArrayList<SlotContent> objectNodes, int prevsize) {
        if (current == null) {
            // First iteration only, initialize the selection.
            current = new HashSet<SchemaCounter>();
            if (schemaNodes != null) {
                SchemaEpisodeNode s = schemaNodes.remove(0);
                current = s.getCountsSet(s.getCounts(level, index));
            } else if (objectNodes != null) {
                SlotContent o = objectNodes.remove(0);
                current = o.getCountsSet(o.getCounts(level, index));
            } else {
                // First iteration, no args? No fun.
                return current;
            }
        }
        if (current.size() == prevsize) {
            // The line has stopped moving, let's see what's going on.
            // Very likely this means that we've hit a long "common" path, and all we need to do is check the remaining nodes.
            Iterator<SchemaCounter> it = current.iterator();
            while (it.hasNext()) {
                SchemaCounter c = it.next();
                if ((!c.episodeNodes.containsAll(schemaNodes)) || (!c.objectNodes.containsAll(objectNodes))) {
                    it.remove();
                }
            }
        }
        if (current.isEmpty()) {
            // Second and later iteration, still no matching counters.
            return current;
        }
        if (schemaNodes.isEmpty()) {
            if (objectNodes.isEmpty()) {
                return current;
            }
            SlotContent o = objectNodes.remove(0);
            current.retainAll(o.getCountsSet(o.getCounts(level, index)));
            return schemaCounterSupersetSelectorOne(level, index, current, schemaNodes, objectNodes);
        } else {
            SchemaEpisodeNode s = schemaNodes.remove(0);
            // Otherwise, chip off the first item, get the respective counters, intersect recursively.
            current.retainAll(s.getCountsSet(s.getCounts(level, index)));
            return schemaCounterSupersetSelectorOne(level, index, current, schemaNodes, objectNodes);
        }
    }
public HashSet<SchemaCounter> schemaCounterSupersetSelectorOne(int level, int index, HashSet<SchemaCounter> current, ArrayList<SchemaEpisodeNode> schemaNodes, ArrayList<SlotContent> objectNodes) {
    // In this particular case, the index,level arguments are dummy.
    return schemaCounterSupersetSelectorOne(level, index, current, schemaNodes, objectNodes, -1);
}


    private int getActualNumberOfVariations(Intention node) {
        if (node.getAssociatedNode() == null) return 0;
        int actionVariations;
        int intentionVariations;
        intentionVariations = 0;
        for (Action action : node.getSubNodes()) {
            actionVariations = 1;
            if (action.getAssociatedNode() == null) continue;
            for (Intention intention : action.getSubNodes()) {
                int temp = getActualNumberOfVariations(intention);
             //   if (temp > 0) {
                    actionVariations *= temp;
             //   }
            }
            intentionVariations += actionVariations;
        }
        return intentionVariations;
    }

    public void generateStatistics(String fileName) {
        try {
            FileWriter outFile = new FileWriter(fileName);
            PrintWriter out = new PrintWriter(outFile);

            Collection<SchemaEpisodeNode> eNodes = new HashSet<SchemaEpisodeNode>();
            Collection<SlotContent> oNodes = new HashSet<SlotContent>();
            int count;
            int sum;

            for (Intention n : mem.getDecisionTree().topLevelGoals.values()) {
                out.print(n.getName() + " | Possible:");
                out.print(n.getPossibleSubTrees() + "| Actual:");
                if (n.getAssociatedNode() == null) {
                    out.println("0");
                    continue;
                } else {
                    out.println(getActualNumberOfVariations(n));
                }
            }
            out.println();

            for (SchemaEpisodeNode n : schemaENodes.values()) {
                //only generate for goals since only there is variability
                if (n.getAssociatedNode().getType() == NodeType.INTENTION) {
                    sum = 0;
                    for (Node a : n.getAssociatedNode().getAllChildrenNodes()) {
                        eNodes.clear();
                        eNodes.add(n);
                        if (a.getAssociatedNode() == null) {
                            continue;
                        }
                        eNodes.add(a.getAssociatedNode());
                        // TODO: Generate statistics for other levels as well.
                        sum += getCount(0, 0, eNodes, oNodes);
                    }
                    out.print(n.getName() + "|");
                    for (Node a : n.getAssociatedNode().getAllChildrenNodes()) {
                        eNodes.clear();
                        eNodes.add(n);
                        if (a.getAssociatedNode() == null) {
                            out.print(a.getName() + " 0 (0%) | ");
                            continue;
                        }
                        eNodes.add(a.getAssociatedNode());
                        // TODO: Generate statistics for other levels as well.
                        count = getCount(0, 0, eNodes, oNodes);
                        out.print(a.getName() + " " + count + " (" + ((double)count*100/sum) + "%) |");
                    }
                    out.println();
                }
            }

            out.flush();
            out.close();
        } catch (Exception e) {
            System.err.println(e.toString());
            System.err.println("Could not create file with statistics!");
        }
    }

    public void performSleep() {
        currentDay += 1;
    }
}
